import os
import sys

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2,4"


project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
print(project_dir_path)
sys.path.append(project_dir_path)

from omics_bert_label_encoding.label.make_labels_charls import MakeLabelsCharls
from omics_bert_label_encoding.label.check_label_files_charls import CheckLabelFilesCharls
from omics_bert_label_encoding.label.processing_labels_charls import process_label_embeddings_charls


def main(dataset_path, save_path, dta_file_path, api_key, instance_mapping, expected_lines, regenerate, model_name, tensor_parallel_size):
    print("--- Step 1: Build Labels ---")
    build_labels = MakeLabelsCharls(dataset_path, save_path, dta_file_path, api_key, instance_mapping)

    build_labels.main()

    print("\n--- Step 2: Check and Regenerate Labels ---")
    check_path = os.path.join(save_path, "labels")

    checker = CheckLabelFilesCharls(
        check_path=check_path,
        expected_line_count=expected_lines,
        regenerate=regenerate,
        dataset_path=dataset_path,
        dta_file_path=dta_file_path,
        api_key=api_key,
        instance_mapping=instance_mapping
    )
    checker.main()

    print("\n--- Step 3: Process Label Embeddings ---")
    labels_dir = os.path.join(save_path, "labels")
    process_label_embeddings_charls(
        labels_json_path=labels_dir,
        output_dir=save_path,
        model_name=model_name,
        tensor_parallel_size=tensor_parallel_size
    )


if __name__ == '__main__':
    project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
    print(project_dir_path)
    sys.path.append(project_dir_path)

    dataset_path = os.path.join(project_dir_path, "datasets/ready/charls/v1/clhls_2000_2018")
    save_path = os.path.join(project_dir_path, "embeddings_local/charls/v1/labels_embeddings/clhls_2000_2018")
    dta_file_path = os.path.join(project_dir_path, "datasets/raw/charls/clhls_2000_2018_longitudinal_dataset_released_version1_modified.dta")

    api_key = ""
    expected_lines = 12
    regenerate = True
    model_name = "Qwen/Qwen3-Embedding-8B"
    tensor_parallel_size = 2

    instance_mapping = {
        "_2": "2002",
        "_5": "2008",
        "_8": "2008",
        "_11": "2011",
        "_14": "2014",
        "_18": "2018",
        None: "baseline"
    }

    main(dataset_path, save_path, dta_file_path, api_key, instance_mapping, expected_lines, regenerate, model_name, tensor_parallel_size)
